from gluoncv.data.transforms import bbox as bbox_transform
from random import randint
from random import random
import gluoncv,cv2
import numpy as np
import mxnet
import mxnet as mx
class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, *args):
        for t in self.transforms:
            args = t(*args)
        return args
    def __len__(self):
        return len(self.transforms)
class bbox_flip():
    def __call__(self,img,bbox,flip_y = False,flip_x = True):
        height,width = img.shape[:2]
        bbox = bbox.copy()
        if flip_y:
            ymax = height - bbox[:, 1]
            ymin = height - bbox[:, 3]
            bbox[:, 1] = ymin
            bbox[:, 3] = ymax
            img= img[::-1, :, :]
        if flip_x:
            xmax = width - bbox[:, 0]
            xmin = width - bbox[:, 2]
            bbox[:, 0] = xmin
            bbox[:, 2] = xmax
            img= img[:, ::-1, :]            
        return img,bbox

class bbox_random_h_flip():
    def __call__(self,img,bbox):
        height,width = img.shape[:2]
        bbox = bbox.copy()
        if randint(0,1):
            xmax = width - bbox[:, 0]
            xmin = width - bbox[:, 2]
            bbox[:, 0] = xmin
            bbox[:, 2] = xmax
            img= img[:, ::-1, :]
        return img,bbox
class bbox_random_rotate():
    def __init__(self,min_angle = -3,max_angle = 3):
        self.min_angle = min_angle
        self.max_angle = max_angle
    def tranform_point(self,matrix,point):
        x,y = point
        p0 = np.array([x, y]).astype(np.float32).reshape(1, 1, 2)
        p = cv2.transform(p0, matrix).reshape(1, -1)
        x_r = p[0, 0]
        y_r = p[0, 1]
        return np.array([x_r,y_r])
    def __call__(self,img,bbox):
        angle = randint(self.min_angle, self.max_angle)
        matrix = cv2.getRotationMatrix2D((img.shape[1] // 2., img.shape[0] / 2.), angle, 1.0);
        img_r = cv2.warpAffine(img, matrix, (img.shape[1], img.shape[0]))
        bbox_r = []
        for i in range(len(bbox)):
            xmin,ymin,xmax,ymax = bbox[i,:4]
            point0 = (xmin,ymin)
            point1 = (xmax,ymin)
            point2 = (xmax,ymax)
            point3 = (xmin,ymax)
            points = [point0,point1,point2,point3]
            points = list(map(lambda x:self.tranform_point(matrix,x),points))
            points = np.array(points)

            xmin = np.min(points[:,0])
            ymin = np.min(points[:,1])
            xmax = np.max(points[:,0])
            ymax = np.max(points[:,1])

            xmin = max(0,xmin)
            ymin = max(0,ymin)
            xmax = min(xmax,img.shape[1]-1)
            ymax = min(ymax,img.shape[0]-1)
            bb = [xmin,ymin,xmax,ymax] + bbox[i,4:].tolist()
            bbox_r.append(bb)
        bbox_r = np.array(bbox_r).conjugate()
        return img_r,bbox_r
class bbox_random_resize():
    def __init__(self,min_scale = 0.8,max_scale = 1.2):
        self.min_scale = 0.7
        self.max_scale = 1.2
    def __call__(self, img,bbox):
        fscale = random() * (self.max_scale-self.min_scale) + self.min_scale
        dshape = [int(fscale * img.shape[0]),int(fscale*img.shape[1])]
        img_r = cv2.resize(img,(dshape[1],dshape[0]))
        bbox_r = bbox_transform.resize(bbox,(img.shape[1],img.shape[0]),(dshape[1],dshape[0]))
        return img_r,bbox_r
class bbox_resize_short():
    def __init__(self,short= 600):
        self._short = short
    def __call__(self, img,bbox):
        if isinstance(img,mx.nd.NDArray):
            img = img.asnumpy()
        if isinstance(bbox,mx.nd.NDArray):
            bbox = bbox.asnumpy()
        fscale = self._short * 1.0 / min(img.shape[:2])
        img_r = cv2.resize(img,(0,0),fx=fscale,fy=fscale)
        bbox_r = bbox_transform.resize(bbox,(img.shape[1],img.shape[0]),(img_r.shape[1],img_r.shape[0]))
        return img_r,bbox_r
class bbox_image_pad():
    def __init__(self,dest_shape = (768,768)):
        self.dest_shape = dest_shape #h,w
    def __call__(self,img_ori,bboxes):
        if isinstance(img_ori,mxnet.nd.NDArray):
            img_ori = img_ori.asnumpy()
        if isinstance(bboxes,mxnet.nd.NDArray):
            bboxes = bboxes.asnumpy()
        if bboxes is not None:
            bboxes = np.array(bboxes)
            assert bboxes.shape[0] >= 1

            dshape = self.dest_shape
            fscale = min(dshape[0]/img_ori.shape[0],dshape[1]/img_ori.shape[1])
            img_resized = cv2.resize(img_ori,dsize=(0,0),fx=fscale,fy=fscale)  # type: numpy.ndarray
            bboxes = bbox_transform.resize(bboxes,(img_ori.shape[1],img_ori.shape[0]),(img_resized.shape[1],img_resized.shape[0]))
            img_padded = np.zeros(shape = (int(dshape[0]),int(dshape[1]),3),dtype=np.float32)
            img_padded[:img_resized.shape[0],:img_resized.shape[1],:img_resized.shape[2]] = img_resized
        return img_padded,bboxes
class bbox_image_pad_n():
    def __init__(self,n=16):
        self.n = n
    def __call__(self,img_ori,bboxes):
        if isinstance(img_ori,mxnet.nd.NDArray):
            img_ori = img_ori.asnumpy()
        if isinstance(bboxes,mxnet.nd.NDArray):
            bboxes = bboxes.asnumpy()
        h,w,c = img_ori.shape
        padn = lambda x:x if x% self.n==0 else x + self.n - x % self.n
        h_,w_ = padn(h),padn(w)
        img_r = np.zeros(shape=(h_,w_,c),dtype=img_ori.dtype)
        img_r[:h,:w,:c] = img_ori
        return img_r,bboxes

class bbox_random_affline():
    def __init__(self):
        pass
    def tranform_point(self,matrix,point):
        x,y = point
        p0 = np.array([x, y]).astype(np.float32).reshape(1, 1, 2)
        p = cv2.perspectiveTransform(p0, matrix).reshape(1, -1)
        x_r = p[0, 0]
        y_r = p[0, 1]
        return np.array([x_r,y_r])
    def __call__(self,img,bbox):
        h = img.shape[0]
        w = img.shape[1]
        ra = randint
        mb = 100
        pts3 = np.float32([[0,0],[w,0],[0,h],[w,h]])
        pts4 = np.float32([[ra(0,mb),ra(0,mb)],[w-ra(0,mb),ra(0,mb)],[ra(0,mb),h-ra(0,mb)],[w-ra(0,mb),h-ra(0,mb)]])
        matrix = cv2.getPerspectiveTransform(pts3,pts4)        
        img_r = cv2.warpPerspective(img, matrix, (img.shape[1], img.shape[0]))
        bbox_r = []
        for i in range(len(bbox)):
            xmin,ymin,xmax,ymax = bbox[i,:4]
            point0 = (xmin,ymin)
            point1 = (xmax,ymin)
            point2 = (xmax,ymax)
            point3 = (xmin,ymax)
            points = [point0,point1,point2,point3]
            points = list(map(lambda x:self.tranform_point(matrix,x),points))
            points = np.array(points)

            xmin = np.min(points[:,0])
            ymin = np.min(points[:,1])
            xmax = np.max(points[:,0])
            ymax = np.max(points[:,1])

            xmin = max(0,xmin)
            ymin = max(0,ymin)
            xmax = min(xmax,img.shape[1]-1)
            ymax = min(ymax,img.shape[0]-1)
            bb = [xmin,ymin,xmax,ymax] + bbox[i,4:].tolist()
            bbox_r.append(bb)
        bbox_r = np.array(bbox_r).conjugate()
        return img_r,bbox_r
class bbox_image_transpose():
    def __call__(self,img,bbox):
        img = np.transpose(img,(2,0,1))
        return img,bbox
class bbox_Aug_train():
    def __init__(self):
        self.img_pad = bbox_image_pad()
        self.bbox_resize = bbox_random_resize()
        self.aug = Compose([self.img_pad])
    def __call__(self,img,bbox):
        return self.aug(img,bbox)
    def __len__(self):
        return len(self.aug)
class bbox_Aug_val():
    def __init__(self):
        pass
    def __call__(self,img,bbox):
        img = img.astype(np.float32)
        bbox = np.array(bbox).astype(np.float32)
        return img,bbox
    def __len__(self):
        return 0